from flask import Flask, request, jsonify, Response ,Blueprint
import pymysql
import chromadb
import json

# 创建蓝图
blue = Blueprint("user",__name__)

# MySQL连接信息
def get_connection(host, user, password, db):
    # 创建数据库连接
    conn = pymysql.connect(
        host=host,
        user=user,
        password=password,
        database=db,
        port=3306
    )
    return conn


# 查询向量数据库并返回结果数组
@blue.route("/query_vdb" , methods=['POST'])
def query_vdb():
    chroma_client = chromadb.HttpClient("192.168.22.101", 8000)
    collection = chroma_client.get_or_create_collection(name="my_collection")

    data = json.loads(request.data)  # 将json字符串转为dict
    value = data['value']
    num = data['num']

    res = collection.query(
        query_texts=value,
        n_results=num
    )
    # 获取documents中的数据，转为一维数组返回
    resArray = [num for row in res.get("documents") for num in row]

    return resArray


@blue.route("/")
def hello_world():
    return "hello_vdb"


@blue.route("/post", methods=['POST'])
def post_world():
    return "hello_vdb"

# 返回json字符串
@blue.route("/json")
def res_json():
    return jsonify({"name":"zhangsan","age":18})
